package cn.liwq.study.demo.java.agent.raw;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

/**
 * <p>Project: demo-common </p>
 * <p>Function:  </p>
 * <p>Description:  </p>
 *
 * @author weiqiang.li
 * @version 1.0
 * @date 2019/2/12
 */
public class TaskTransformer implements ClassFileTransformer {

    private String targetClassName;

    private ClassLoader targetClassLoader;

    public TaskTransformer(String clazzName, ClassLoader classLoader) {
        this.targetClassName = clazzName;
        this.targetClassLoader = classLoader;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        byte[] bytecode = classfileBuffer;

        String finalTargetClassName = this.targetClassName.replaceAll("\\.", "/");

        if (!className.equals(finalTargetClassName)) {
            return bytecode;
        }

        System.out.println(className);

        if (loader.equals(this.targetClassLoader)) {
            System.out.println("[Agent] Transforming class MyTask");
            try {
                ClassPool cp = ClassPool.getDefault();
                CtClass cc = cp.get(targetClassName);
                CtMethod m = cc.getDeclaredMethod("output");
                m.addLocalVariable("startTime", CtClass.longType);
                m.insertBefore("startTime = System.currentTimeMillis();");

                StringBuilder endBlock = new StringBuilder();
                m.addLocalVariable("endTime", CtClass.longType);
                m.addLocalVariable("opTime", CtClass.longType);
                endBlock.append("endTime = System.currentTimeMillis();");
                endBlock.append("opTime = (endTime - startTime) / 1000;");

                endBlock.append("System.out.println(\"output waste time: \" + opTime);");

                m.insertAfter(endBlock.toString());

                bytecode = cc.toBytecode();
                cc.detach();
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        }

        return bytecode;
    }
}
